from ctypes import pythonapi
from operator import itemgetter

from langchain.chains import create_sql_query_chain
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_community.vectorstores import FAISS
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import create_retriever_tool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
import re
import ast
from typing import Optional

from langgraph.prebuilt import create_react_agent

db = SQLDatabase.from_uri("sqlite:///Chinook.db")
print(db.dialect)
print(db.get_usable_table_names())
print(db.run("SELECT * FROM Artist LIMIT 10;"))

llm = ChatOpenAI(
    api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
    model='deepseek-ai/DeepSeek-V3',
    base_url="https://www.henapi.top/v1"
)


def extract_executable_sql(llm_output: str) -> Optional[str]:
    """
    从LLM输出中提取可直接执行的SQL语句
    返回: 纯净的SQL语句或None(如果提取失败)
    """
    # 情况1：处理 ```sql ... ``` 格式
    sql_match = re.search(r'```sql\n(.*?)\n```', llm_output, re.DOTALL)
    if sql_match:
        return sql_match.group(1).strip()

    # 情况2：处理 SQLQuery: ... 格式
    prefix_match = re.search(r'SQLQuery:\s*(.*?;)\s*$', llm_output, re.IGNORECASE | re.DOTALL)
    if prefix_match:
        return prefix_match.group(1).strip()

    # 情况3：尝试提取标准的SELECT/INSERT/UPDATE等语句
    sql_match = re.search(
        r'\b(SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP)\b.*?;',
        llm_output,
        re.IGNORECASE | re.DOTALL
    )
    if sql_match:
        return sql_match.group(0).strip()

    return None


# chain=create_sql_query_chain(llm,db)
# response=chain.invoke({"question": "How many employees are there"})
# print("llm生成的结果："+response)
# response=extract_executable_sql(response)
# print(response)

# print(db.run(response))


execute_query = QuerySQLDataBaseTool(db=db)
write_query = create_sql_query_chain(llm, db)
# chain= write_query|extract_executable_sql|execute_query
# resp=chain.invoke({"question": "How many employees are there"})
# print(resp)

answer_prompt = PromptTemplate.from_template(
    """Given the following user question, corresponding SQL query, and SQL result, answer the user question.

Question: {question}
SQL Query: {query}
SQL Result: {result}
Answer: """
)

chain = (
        RunnablePassthrough.assign(query=write_query|extract_executable_sql).assign(
            result=itemgetter("query")  | execute_query)
        | answer_prompt
        | llm
        | StrOutputParser()
)
# resp = chain.invoke({"question": "How many employees are there"})
# print(resp)


toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()
print(tools)

SQL_PREFIX = """You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

To start you should ALWAYS look at the tables in the database to see what you can query.
Do NOT skip this step.
Then you should query the schema of the most relevant tables."""

system_message = SystemMessage(content=SQL_PREFIX)

agent_executor = create_react_agent(llm, tools, prompt=system_message)
# for s in agent_executor.stream(
#     {"messages": [HumanMessage(content="Which country's customers spent the most?")]}
# ):
#     print(s)
#     print("----")

# for s in agent_executor.stream(
#     {"messages": [HumanMessage(content="Describe the playlisttrack table")]}
# ):
#     print(s)
#     print("----")

# 输入：数据库连接对象 db 和 SQL 查询字符串 query。
# 步骤：
# 使用 db.run(query) 执行 SQL 查询，获取结果 res。
# 使用 ast.literal_eval(res) 将结果字符串解析为 Python 数据结构（假设结果是一个字符串表示的列表）。
# 使用列表推导式 [el for sub in ast.literal_eval(res) for el in sub if el] 将嵌套列表展平，并过滤掉空元素。
# 使用正则表达式 re.sub(r"\b\d+\b", "", string).strip() 去除字符串中的数字，并去除首尾空格。
# 使用 list(set(res)) 去除重复元素，返回一个唯一值列表。
# 输出：处理后的唯一值列表。

def query_as_list(db, query):
    res = db.run(query)
    res = [el for sub in ast.literal_eval(res) for el in sub if el]
    res = [re.sub(r"\b\d+\b", "", string).strip() for string in res]
    return list(set(res))


artists = query_as_list(db, "SELECT Name FROM Artist")
albums = query_as_list(db, "SELECT Title FROM Album")
# print(albums[:5])
openAIEmbeddings=OpenAIEmbeddings(api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
    model='text-embedding-3-small',
    base_url="https://www.henapi.top/v1")

vector_db = FAISS.from_texts(artists + albums, openAIEmbeddings)
retriever = vector_db.as_retriever(search_kwargs={"k": 5})
description = """Use to look up values to filter on. Input is an approximate spelling of the proper noun, output is \
valid proper nouns. Use the noun most similar to the search."""
retriever_tool = create_retriever_tool(
    retriever,
    name="search_proper_nouns",
    description=description,
)

print(retriever_tool.invoke("Alice Chains"))

# python3.9装了faiss-cpu 1.8